import copy, os, importlib
from multiprocessing import Pool
from itertools import repeat

import yaml
import numpy as np
from tqdm import tqdm

from .basic_planner import basic_planner
from ..minimization_variables import basic_variable
from ..target_list import target_list
from ..satellite_formation import basic_sate_form
from ..file_io.output_control import global_output

class ant_result():
    def __init__(self, total_time, delta_angle, trace, ant_path):
        self.total_time = total_time
        self.delta_angle = delta_angle
        self.trace = trace
        self.ant_path = ant_path
class acs_ant():
    def __init__(self, target_list: target_list, sate_form: basic_sate_form, min_var: basic_variable, pheromone_matrix: np.ndarray, yaml_config: yaml.Node):
        self.__target_list = target_list
        self.__sate_form = sate_form
        self.__minimizi_var = min_var
        self.__pheromone_matrix = pheromone_matrix

        self.__visual_angle = self.__sate_form.visual_angle
        self.__sleep_time   = self.__sate_form.sleep_time          # 无观测目标时的等待时间 day
        self.__alpha        = yaml_config.get("TrailWeight", 1.)        # 成本权重 \alpha
        self.__beta         = yaml_config.get("HeuristicWeight", 1.)    # 启发式因子权重 \beta
        self.__q0           = yaml_config.get("ExplorationFactor", 0.1) # Exploration因子 q_0

        self.__total_time = 0.
        self.__n_sleep = 0.
        self.__delta_angle = 0.
        self.__trace = np.zeros((self.__target_list.len, self.__target_list.len), dtype=np.int8)
        self.__ant_path = []

    @property
    def trace(self):
        return self.__trace

    @property
    def n_sleep(self):
        return self.__n_sleep

    @property
    def ant_path(self):
        return self.__ant_path

    @property
    def total_time(self):
        return self.__total_time

    @property
    def sate_form(self):
        return self.__sate_form

    def init_targets(self):
        while True:
            self.__target_list_iter = self.__minimizi_var.filter(self.__target_list.nplist, self.__sate_form)
            #self.filter()

            if len(self.__target_list_iter) > 0:
                break
            # 当前视角下无可观测任务
            else:
                self.__total_time = self.__total_time + self.__sleep_time
                self.__sate_form.update_satellite(sleep = True)

        return self.__target_list_iter, self.__total_time

    def set_init_target(self, init_target):
        # 更新队列信息
        delta_lon = np.abs(init_target[1]-self.__sate_form.point_lon) % 360
        #target_retarget_time = delta_lon / self.__sate_form.angular_velocity
        target_finish_time = self.__minimizi_var.evaluate(np.array([init_target]), self.__sate_form)[0]
        init_target_id, time_cost = self.__sate_form.update_satellite(init_target, target_finish_time)

        self.__total_time = target_finish_time
        self.__delta_angle = delta_lon
        self.__target_list.delete_target(init_target_id)
        self.__ant_path.append([init_target_id, time_cost, delta_lon, 0.])

        #print(f"{len(self.__target_list_iter)}, {init_target_id}")
        return [init_target_id, time_cost, delta_lon, 0.]

    def execute(self):
        while self.__target_list.len > 0:
            # 找到下一个观测目标
            point_angle = self.__sate_form.point_lon
            #self.__target_list_iter = self.__target_list.nplist
            self.__target_list_iter = self.__minimizi_var.filter(self.__target_list.nplist, self.__sate_form)
            #self.filter()
            target_id, time_cost = self.select(self.__ant_path[-1][0])
            self.__total_time = self.__total_time + time_cost
            self.__delta_angle = self.__delta_angle + (np.abs(point_angle - self.__sate_form.point_lon) % 360)

            # 如无合适目标，则等待一段时间后继续
            if target_id == -1:
                self.__n_sleep = self.__n_sleep + 1
                continue

            # 标记蚂蚁路径
            self.__trace[int(self.__ant_path[-1][0] - 1), int(target_id - 1)] = 1
            self.__ant_path.append([target_id, self.__total_time, self.__delta_angle, self.__n_sleep])

            # 更新信息
            self.__target_list.delete_target(target_id)
            self.__n_sleep = 0

            if self.__n_sleep * self.__sleep_time > 400.:
                print("It seems that some tasks are unable to be visited with current constraints.")
                print("BREAK! Remainning tasks:")
                print("ID,\tra,\t,dec,\tt_tot")
                for task in self.__target_list:
                    print(f"{task[0]},\t{task[1]},\t{task[2]},\t{task[3]}")
                break

        return self.__total_time, self.__delta_angle

    def select(self, last_target_id: int):
        # 如果无满足条件的待测目标，则等待一段时间
        if(len(self.__target_list_iter) == 0):
            #global_output.write_line(-1)
            return self.__sate_form.update_satellite(sleep=True)

        # 计算剩余目标重定向需要的时间
        #target_retarget_time = (self.__target_list_iter[:,1]-self.__sate_form.point_lon) / self.__sate_form.angular_velocity
        #delta_lon = np.abs(self.__target_list_iter[:,1]-self.__sate_form.point_lon) % 360
        #target_retarget_time = delta_lon / self.__sate_form.angular_velocity
        target_finish_time = self.__minimizi_var.evaluate(self.__target_list_iter, self.__sate_form)

        # 根据蚁群算法确定目标
        eta_matrix = 1. / target_finish_time
        #eta_matrix = 1. / (target_retarget_time + self.__target_list_iter[:,3])
        target_index = np.array(self.__target_list_iter[:,0], dtype=int)
        tau_matrix = self.__pheromone_matrix[int(last_target_id - 1)][target_index - 1]
        choose_prob = np.power(eta_matrix, self.__alpha) * np.power(tau_matrix, self.__beta)
        choose_prob = choose_prob / np.sum(choose_prob)

        q_exploration = np.random.random()
        if q_exploration < self.__q0:
            index = np.argmax(choose_prob)
        else:
            index = np.random.choice(range(len(target_index)), size=1, p=choose_prob)

        index = int(index)
        current_target = self.__target_list_iter[index]

        #print(type(target_finish_time[index]))
        finish_time = target_finish_time[index]
        return self.__sate_form.update_satellite(current_target, finish_time)

    def update_satellite(self, current_target=[], total_time_cost = 0., sleep = False):
        if sleep == True:
            self.__sate_form.orbit_lon = self.__sate_form.orbit_lon + self.__sleep_time * self.__sate_form.orbit_rotate_angular_velocity
            return -1, self.__sleep_time

        # 更新卫星编队状态
        self.__sate_form.point_lon = current_target[1]
        self.__sate_form.point_lat = current_target[2]
        self.__sate_form.orbit_lon = self.__sate_form.orbit_lon + total_time_cost * self.__sate_form.orbit_rotate_angular_velocity

        return current_target[0], total_time_cost

class acs_planner(basic_planner):
    '''
    基本的蚁群算法
    '''
    def __init__(self, yaml_config):
        self.__target_list = None
        self.__sate_form = None

        self.__yaml_config      = yaml_config
        self.__tau0             = yaml_config.get("InitialPheromone", 0.1)      # 初始信息素    \tau_0
        self.__phe_gamma        = yaml_config.get("PheromoneUpdateWeight", 1.) # 信息素更新权重 \gamma
        self.__rho              = yaml_config.get("PheromoneVolatilize", 0.8)   # 信息素挥发    \rho
        self.__num_of_ants      = yaml_config.get("NumOfAnts", 30)              # 蚂蚁数目
        self.__iter_num         = int(yaml_config.get("IterNum", 2000))         # 迭代次数
        self.__multi_process    = int(yaml_config.get("MultiProcessing", 10))   # 多进程选项，<=1表示不开启多进程
        self.__global_best      = yaml_config.get("EnableGlobalBest", False)    # 是否使用全局最优蚂蚁更新信息素
        self.__min_var          = yaml_config.get("MinimizationVariable", "simple_time")
        self.__min_var_config   = yaml_config.get("MinimizationConfig", {})

        # 指定最小化变量
        minimize_mod = importlib.import_module(f".minimization_variables.{self.__min_var}", package="miyinff")
        self.__minimize_clz = getattr(minimize_mod, self.__min_var)
        self.__minimize_obj = self.__minimize_clz(self.__min_var_config)

        self.__csv_task     = "acs_planner_result"
        self.__csv_results  = "acs_planner_operation"

        if os.cpu_count() < self.__multi_process:
            print("Warning! Input MultiProcessing larger than number of cpu cores. ")
            print(f"MultiProcessing: {self.__multi_process}, \tCPU cores: {os.cpu_count()}")
            print("Reset MultiProcessing to current core number!")
            self.__multi_process = os.cpu_count()

    def param_list(self):
        return self.__minimize_clz.param_list()

    #def sate_form_type(self):
    #    return self.__minimize_clz.sate_form_type()

    def initialize(self, target_list, random_seed, **kwargs):

        #self.__target_list = copy.deepcopy(target_list)
        #self.__sate_form = copy.deepcopy(sate_format)
        self.__target_list, self.__sate_form = self.__minimize_obj.initialize(target_list)
        self.__random_seed = random_seed

        self.__pheromone_matrix = np.full((self.__target_list.len, self.__target_list.len), self.__tau0, dtype=np.float32)

        init_ant = acs_ant(copy.deepcopy(self.__target_list), copy.deepcopy(self.__sate_form), copy.deepcopy(self.__minimize_obj), \
                           self.__pheromone_matrix, self.__yaml_config)
        self.__init_targets, self.__init_sleep_time = init_ant.init_targets()

        if self.__init_sleep_time > 0.:
            global_output.write_log(f"Warning! Non of targets can be detected at initial state. Sleep {self.__init_sleep_time} days.")

        #self.__init_target_ids = self.__init_targets[:,0]
        self.__init_target_number = len(self.__init_targets)
        global_output.write_log(f"There are {self.__init_target_number} targets can be detected at initial state.")
        if self.__num_of_ants < self.__init_target_number:
            global_output.write_log(f"Warning! Request number of ants '{self.__num_of_ants}' less than number of initial targets '{self.__init_target_number}'")
            global_output.write_log(f"Reset number of ants to '{self.__init_target_number}'")

        self.__best_ant = None
        self.__iter_best_time = []

    def operation(self):
        for iter in tqdm(range(self.__iter_num)):
            # 一次迭代
            results = []
            if self.__multi_process > 1:
                pool = Pool(processes=self.__multi_process)
                results = pool.starmap(self.process_ant, list(zip(repeat((self), self.__num_of_ants), repeat(iter, self.__num_of_ants), range(self.__num_of_ants))))
                #results = pool.starmap(self.process_ant, list(repeat((self,), self.__num_of_ants)))
                pool.close()
                pool.join()

            else:
                for i_ant in range(self.__num_of_ants):
                    results.append(self.process_ant(self, iter, i_ant))

            # 单次循环最佳蚂蚁
            ant_times = np.array([res.total_time for res in results], dtype=np.float32)
            i_best = np.argmin(ant_times)
            iter_best_ant = results[i_best]

            # 计算信息素更新权重
            if self.__best_ant == None:
                self.__best_ant = iter_best_ant
            # 损失函数
            # sigmoid function: 1 / ( 1 + e^{-X})
            #phe_weight = 1. / (1 + np.power(np.e, iter_best_ant.total_time - self.__best_ant.total_time))

            # logistic function: ln(1 + e^{X})
            phe_weight = np.log(1 + np.power(np.e, self.__best_ant.total_time - iter_best_ant.total_time))

            # 倒数
            #phe_weight = 1.
            #if int(self.__best_ant.total_time) < int(iter_best_ant.total_time):
            #    phe_weight = 1. / (int(iter_best_ant.total_time) - int(self.__best_ant.total_time))

            # 更新信息素
            if self.__global_best:
                self.__pheromone_matrix = self.__pheromone_matrix * self.__rho + \
                    np.power(phe_weight, self.__phe_gamma) * self.__best_ant.trace
            else:
                self.__pheromone_matrix = self.__pheromone_matrix * self.__rho + \
                    np.power(phe_weight, self.__phe_gamma) * iter_best_ant.trace
                    #(1-self.__rho) * phe_weight * iter_best_ant.trace

            # 记录信息素矩阵变化
            global_output.write_log(f"Iter: {iter}, Min: {np.min(self.__pheromone_matrix)}, Max: {np.max(self.__pheromone_matrix)}, Sum: {np.sum(self.__pheromone_matrix)}", print_out=False)

            # 记录全局最佳蚂蚁
            if self.__best_ant.total_time > iter_best_ant.total_time:
                self.__best_ant = iter_best_ant
            self.__iter_best_time.append(iter_best_ant.total_time)


        # 输出最优蚂蚁路线
        output_csv = global_output.create_csv_file(self.__csv_task)
        output_csv.write_head("ID")
        output_csv.write_head(self.__target_list.param_list)
        output_csv.write_head(["Time", "Delta_angle", "N_sleep"])

        for path in self.__best_ant.ant_path:
            target = self.__target_list.find_target(path[0])
            line = target + path[1:]
            output_csv.write_line(line)
        output_csv.save()
        self.__sate_form.end_operation()
        global_output.write_log(f"\nOperation finished, Best Time: {self.__best_ant.total_time}, Delta angle: {self.__best_ant.delta_angle}")

        # 输出优化过程
        planner_csv = global_output.create_csv_file(self.__csv_results)
        planner_csv.write_head(["I", "Time"])
        for i, time in enumerate(self.__iter_best_time):
            planner_csv.write_line([i, time])
        planner_csv.save()

    @staticmethod
    def process_ant(self, i_iter, i_ant):
        # 对于多线程，需重新设置随机数
        np.random.seed(int(self.__random_seed + i_iter*1000 + i_ant))
        a_ant = acs_ant(target_list=copy.deepcopy(self.__target_list), sate_form=copy.deepcopy(self.__sate_form), min_var=self.__minimize_obj, \
                        pheromone_matrix=self.__pheromone_matrix, yaml_config=self.__yaml_config)

        if i_ant < self.__init_target_number:
            a_ant.set_init_target(self.__init_targets[i_ant])
        else:
            a_ant.set_init_target(self.__init_targets[int(np.random.random() * self.__init_target_number)])

        total_time, delta_angle = a_ant.execute()

        return ant_result(total_time, delta_angle, a_ant.trace, a_ant.ant_path)

    def clear(self):
        pass
